import argparse
from pathlib import Path
from typing import ClassVar
import polars as pl

from pre_base import DATA_DIR, PrepareBase


class PrepareClingen(PrepareBase):
    name: ClassVar[str] = "DECIPHER"
    header: ClassVar[str] = '##INFO=<ID=DECIPHER,Number=1,Type=String,Description="DECIPHER, accession:chrom:start:end:gain_af:gain_loss">'

    def read_input(self):
        dataframe = pl.read_csv(
            self.input_file,
            separator="\t",
            columns=["#population_cnv_id", "chr", "start", "end", "deletion_observations", "duplication_observations", "observations"],
            schema_overrides={"chr": str},
        )
        dataframe = (
            dataframe.with_columns(
                pl.format("chr{}", pl.col("chr")).alias("chr"),
                (pl.col("duplication_observations") / pl.col("observations")).alias("gain_af"),
                (pl.col("deletion_observations") / pl.col("observations")).alias("loss_af"),
            )
            .with_columns(
                pl.format(
                    "{}:{}:{}:{}:{}:{}",
                    pl.col("#population_cnv_id"),
                    pl.col("chr"),
                    pl.col("start"),
                    pl.col("end"),
                    pl.col("gain_af"),
                    pl.col("loss_af"),
                ).alias("INFO")
            )
            .select(["chr", "start", "end", "INFO"])
            .sort(["chr", "start", "end"])
        )
        return dataframe


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Prepare DECIPHER")
    parser.add_argument("--input_file", "-i", default=DATA_DIR / "raw" / "decipher_population_cnv_grch38.txt.gz", type=Path, help="input file")
    parser.add_argument("--outdir", "-o", default=DATA_DIR / "bed", type=Path, help="output directory")
    parser.set_defaults(func=lambda args: PrepareClingen.model_validate(vars(args)).run())
    args = parser.parse_args()
    args.func(args)

    # f"{variant_id}:{self.record.chrom}:{start}:{end}:{revstat}:{clnsig}"
